from District import district
from Intake import intake
from River import river
from Station import station
from Water import waterDemand
from CalculateUtilService import calculateUtilService
class wRASimulationService:
    def __init__(self):
        self.time = None

    def execute(self, river: river, T: int, startT: int, priority_supply : list,priority_demand:list):
        # 从水文年开始计算
        self.time = [(i + startT) % T for i in range(T)]
        #初始化水位
        for i in range(len(river.getStations())):
            river.getStations()[i].getTimeBeans()[startT].setvBegin(river.getStations()[i].getvMax())
        for priority in priority_supply:
            if priority == '外调水':
                # 外调水分配
                for i in range(len(river.getDistricts())):
                    d = river.getDistricts()[i]
                    QNSBD = d.getwWDS() * 1e8 / (T * 24 * 3600)  # 外调水
                    for t in range(T):
                        lack = d.getAw()[t] + d.getDw()[t] + d.getIw()[t] + d.getEw()[t]
                        Q = min(lack, QNSBD)
                        self.allocation(river.getDistricts()[i], t, Q,priority_demand)
            elif priority == '地下水':
                # 地下水水分配
                for i in range(len(river.getDistricts())):
                    d = river.getDistricts()[i]
                    QNSBD = d.getGwMax() * 1e8 / (T * 24 * 3600)  # 地下水
                    for t in range(T):
                        lack = d.getAw()[t] + d.getDw()[t] + d.getIw()[t] + d.getEw()[t]
                        Q = min(lack, QNSBD)
                        self.allocation(river.getDistricts()[i], t, Q,priority_demand)
            elif priority == '本地水':
                for i in range(T):
                    self.bendifenshui(river, i,priority_demand)
            elif priority == '水库水':
                for i in range(T):
                    self.shuikufenshui(river, i,priority_demand)

    def bendifenshui(self, River: river, t: int,priority_demand:list) -> float:
        for i in range(len(River.getElements())):
            element = River.getElements()[i]
            if isinstance(element, river):
                r = element
                flowInward = self.bendifenshui(r, self.time[t],priority_demand)
                if not r.getEnd():
                    inflowIntvl = River.getElements()[i + 1].getTimeBeans()[self.time[t]].getInflowIntvl()
                    inflow = r.getTimeBeans()[self.time[t]].getInflowUpper() + inflowIntvl + flowInward
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowIntvl(
                        inflowIntvl + flowInward)
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflow(inflow)
                    outflow = inflow
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setOutflow(outflow)
                else:
                    return flowInward
            elif isinstance(element, station):
                s = element
                inflow = s.getTimeBeans()[self.time[t]].getInflowIntvl() + s.getTimeBeans()[
                    self.time[t]].getInflowUpper()
                s.getTimeBeans()[self.time[t]].setInflow(inflow)
                outflow =inflow
                if not s.getEnd():
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowUpper(outflow)
                else:
                    return outflow
            elif isinstance(element, intake):
                take = element
                inflow = take.getTimeBeans()[self.time[t]].getInflowIntvl() + take.getTimeBeans()[
                    self.time[t]].getInflowUpper()
                take.getTimeBeans()[self.time[t]].setInflow(inflow)

                demand = take.getDistrict().getAw()[self.time[t]] + take.getDistrict().getIw()[
                        self.time[t]] + \
                         take.getDistrict().getDw()[self.time[t]]
                capacity = take.getCapacity()
                available = inflow - take.getOutflowMin()
                if available < 0:
                    available = 0
                flowWdrl = min(min(demand, capacity), available)
                take.getTimeBeans()[self.time[t]].setFlowWdrl(flowWdrl)
                outflow = inflow - flowWdrl
                take.getTimeBeans()[self.time[t]].setOutflow(outflow)
                self.allocation(take.getDistrict(), self.time[t],flowWdrl,priority_demand)
                if not take.getEnd():
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowUpper(outflow)
                else:
                    return outflow
        return -1

    def shuikufenshui(self, River: river, t: int,priority_demand:list) -> float:
        c=calculateUtilService()
        for i in range(len(River.getElements())):
            element = River.getElements()[i]
            if isinstance(element, river):
                r = element
                flowInward = self.shuikufenshui(r, self.time[t],priority_demand)
                if not r.getEnd():
                    inflowIntvl = River.getElements()[i + 1].getTimeBeans()[self.time[t]].getInflowIntvl()
                    inflow = r.getTimeBeans()[self.time[t]].getInflowUpper() + inflowIntvl + flowInward
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowIntvl(
                            inflowIntvl + flowInward)
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflow(inflow)
                    outflow = inflow
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setOutflow(outflow)
                else:
                    return flowInward
            elif isinstance(element, station):
                s = element
                flow = s.getTimeBeans()[self.time[t]].getInflowIntvl()
                inflow = s.getTimeBeans()[self.time[t]].getInflowIntvl() + s.getTimeBeans()[
                    self.time[t]].getInflowUpper()
                s.getTimeBeans()[self.time[t]].setInflow(inflow)
                vLimit = c.getvMax(s, self.time[t])
                demand = s.getDistrict().getAw()[self.time[t]] + s.getDistrict().getIw()[self.time[t]] + \
                         s.getDistrict().getDw()[self.time[t]]
                capacity = s.getCapacity()
                available = c.getQavl(inflow,s.getTimeBeans()[self.time[t]].getvBegin(),s.getOutflowMin(), s.getvMin())
                if available < 0:
                    available = 0
                flowWdrl = min(min(demand, capacity), available)
                s.getTimeBeans()[self.time[t]].setFlowWdrl(flowWdrl)

                outflow = s.getOutflowMin()
                vEnd = c.getV1(flow, outflow, flowWdrl,
                                                           s.getTimeBeans()[self.time[t]].getvBegin())
                if vEnd > vLimit:
                    outflow = calculateUtilService.getQout(flow, flowWdrl,
                                                                    s.getTimeBeans()[self.time[t]].getvBegin(),
                                                                    vLimit)
                s.getTimeBeans()[self.time[t]].setOutflow(outflow)

                self.allocation(s.getDistrict(), self.time[t],flowWdrl,priority_demand)

                if not s.getEnd():
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowUpper(outflow)
                else:
                    return outflow
            elif isinstance(element, intake):
                take = element
                inflow = take.getTimeBeans()[self.time[t]].getInflowIntvl() + take.getTimeBeans()[
                        self.time[t]].getInflowUpper()
                take.getTimeBeans()[self.time[t]].setInflow(inflow)
                outflow = inflow
                if not take.getEnd():
                    River.getElements()[i + 1].getTimeBeans()[self.time[t]].setInflowUpper(outflow)
                else:
                    return outflow
        return -1


    def allocation(self,d: district, t: int,flowWdrl:float,priority_demand:list ):

        # supply_data = [
        #     WaterSupply('a', 20, 3),
        #     WaterSupply('b', 30, 1),
        #     WaterSupply('c', 20, 4),
        #     WaterSupply('waidiao', flowWdrl, 2),
        # ]

        demand_data = [
            waterDemand('aw', d.getAw()[t], priority_demand[0]),
            waterDemand('iw', d.getIw()[t], priority_demand[1]),
            waterDemand('dw', d.getDw()[t], priority_demand[2]),
            waterDemand('ew', d.getEw()[t], priority_demand[3])
        ]
        total_demand = sum(demand.value for demand in demand_data)
        demand_data.sort(key=lambda x: x.priority, reverse=True)
        if (self.isLack(total_demand)):
            return
        for demand in demand_data:
            if demand.priority != 0:
                allocated_amount=min(demand.value,flowWdrl)
                demand.value -= allocated_amount
                flowWdrl -= allocated_amount
            if flowWdrl == 0:
                break

        total_remaining_demand = sum(demand.value for demand in demand_data)
        if flowWdrl>=total_remaining_demand:
            flowWdrl = flowWdrl - total_remaining_demand
            for demand in demand_data:
                if demand.priority == 0:
                    demand.value = 0
        else:
            for demand in demand_data:
                if demand.priority == 0:
                    demand.value = demand.value - (flowWdrl * (demand.value / total_remaining_demand))

        for demand in demand_data:
            if demand.name == 'aw':
                d.getAw()[t]=demand.value
            elif demand.name == 'iw':
                d.getIw()[t]=demand.value
            elif demand.name == 'dw':
                d.getDw()[t]=demand.value
            elif demand.name == 'ew':
                d.getEw()[t]=demand.value

        # supply_data.sort(key=lambda x: x.priority, reverse=True)
        # total_remaining_supply = sum(supply.value for supply in supply_data if supply.value is not None)
        #
        # for demand in demand_data:
        #     if demand.priority is not None:
        #         for supply in supply_data:
        #             if supply.value is not None:
        #                 allocated_amount = min(demand.value, supply.value)
        #                 demand.value -= allocated_amount
        #                 supply.value -= allocated_amount
        #                 total_remaining_supply -= allocated_amount
        #                 if demand.value == 0:
        #                     break
        #
        # total_remaining_demand = sum(demand.value for demand in demand_data)
        #
        # if total_remaining_supply >= total_remaining_demand:
        #     total_remaining_supply = total_remaining_supply - total_remaining_demand
        #     for demand in demand_data:
        #         if demand.priority is None:
        #             demand.value = 0
        #     for supply in supply_data:
        #         if supply.value is not None and total_remaining_demand > 0:
        #             dec = min(supply.value, total_remaining_demand)
        #             supply.value -= dec
        #             total_remaining_demand -= dec
        #         if total_remaining_demand == 0:
        #             break
        # else:
        #     for demand in demand_data:
        #         if demand.priority is None:
        #             demand.value = demand.value - (total_remaining_supply * (demand.value / total_remaining_demand))
        #     for supply in supply_data:
        #         supply.value = 0
        #     total_remaining_demand = total_remaining_demand - total_remaining_supply
        #     total_remaining_supply = 0

    def isLack(self, lack: float) -> bool:
        return lack > 0.00001
